#include "common.h"
#include <vector>
#include <map>
#include <set>
#include <assert.h>
// #define _VERBOSE
struct Router all_router[MAXROUTERS];
struct Router *this_router;
int n_router;
int myid;
int dv[MAXROUTERS][MAXROUTERS];
int weight[MAXROUTERS];
struct RTE routing_table[MAXROUTERS];

int sockfd[MAXROUTERS], maxfd; // sockfd[myid] refers to the socket to the agent
bool conn[MAXROUTERS];

int SEQ = 0;

struct PropMsg
{
    static int wait_cnt;
    std::set<int> waitlist;
    int father_id;
    int father_seq;
    PropMsg() { father_id = father_seq = -1; }
    PropMsg(int fid, int fseq)
    {
        father_id = fid;
        father_seq = fseq;
        for (int i = 0; i < n_router; ++i)
            if (i != myid && weight[i] < INF)
            {
                wait_cnt++;
                waitlist.insert(i);
            }
    }
    bool del(int id)
    {
        bool ret = waitlist.erase(id);
        wait_cnt -= ret;
        return ret; // 0 failed, 1 success
    }
};
std::map<int, PropMsg> M; // map from seq of msg to struct PropMsg
int PropMsg::wait_cnt = 0;

// void Send(int fd, void *message, size_t m_len, int flags)
// {
//     if (rio_readn(fd, message, m_len) < 0)
//     {
//         perror("Send failed");
//         exit(EXIT_FAILURE);
//     }
// }

// int Recv(int fd, MessageR2R *message)
// {
//     char buf[MAXLINE];
//     int recv_bytes = recv(fd, (void *)buf, MAXLINE, 0);
//     if (recv_bytes < 0)
//         return recv_bytes;
//     Message *msg = (Message *)buf;
//     if (msg->message_type != R2R)
//         return -1;
//     memcpy(message, buf, sizeof(MessageR2R));
//     return recv_bytes;
// }

void get_table(char *str)
{
    str[0] = '\0';
    for (int i = 0; i < n_router; ++i)
        if (routing_table[i].cost < INF)
            routing_table[i].tostring(str + strlen(str));
}
void print_table()
{
    char str[MAXLINE];
    get_table(str);
    puts(str);
}
void print_dv(int row)
{
    printf("[DEBUG]: distance vector of router %d:\n", row);
    for (int i = 0; i < n_router; ++i)
        printf("%d ", dv[row][i]);
    puts("");
}
bool BF()
{
#ifdef _VERBOSE
    puts("[INFO]: B-F");
    puts("[DEBUG]: dv before B-F:");
    print_dv(myid);
#endif
    bool flag = 0;
    for (int i = 0; i < n_router; ++i)
        if (i != myid)
        {
            int next = -1, cost = INF;
            for (int j = 0; j < n_router; ++j)
                if (j != myid && weight[j] < INF)
                    if (dv[j][i] + weight[j] < cost)
                        next = j, cost = dv[j][i] + weight[j];
            routing_table[i].next = next;
            routing_table[i].cost = cost;
            if (dv[myid][i] != cost)
            {
                dv[myid][i] = cost;
                flag = 1;
            }
        }
#ifdef _VERBOSE
    puts("[DEBUG]: dv after B-F:");
    print_dv(myid);
    puts("[DEBUG]: router table after B-F:");
    print_table();
#endif
    return flag;
}

void set_neighbor_dv(int id, int a[])
{
    for (int i = 0; i < n_router; ++i)
        dv[id][i] = a[i];
}
void print_wait()
{
    puts("[INFO]: wait:");
    printf("wait_cnt: %d\n", PropMsg::wait_cnt);
    for (auto s : M)
    {
        printf("message %d waiting for reply from ", s.first);
        for (auto e : s.second.waitlist)
            printf("%d ", e);
        printf("(fa_id: %d, fa_seq: %d)\n", s.second.father_id, s.second.father_seq);
        puts("");
    }
}
void propogate()
{
#ifdef _VERBOSE
    puts("[INFO]: propogate");
#endif

    ++SEQ;
    PropMsg ___(-1, -1);
    if (!___.waitlist.empty())
        M[SEQ] = ___;
    for (int i = 0; i < n_router; ++i)
        if (i != myid && weight[i] < INF)
        {
            MessageR2R msgr2r(SEND, myid, SEQ);
            memcpy(msgr2r.dv, dv[myid], sizeof(msgr2r.dv));
            Rio_writen(sockfd[i], &msgr2r, msgr2r.len);
#ifdef _VERBOSE
            printf("[DEBUG]: Send signal to router %d (seq: %d).\n", i, SEQ);
#endif
        }

    fd_set read_set, ready_set;
    FD_ZERO(&read_set);
    for (int i = 0; i < n_router; ++i)
        if (i != myid)
            FD_SET(sockfd[i], &read_set);

    while (!M.empty())
    {
#ifdef _VERBOSE
        print_wait();
#endif

        ready_set = read_set;
        while (select(maxfd, &ready_set, NULL, NULL, NULL) <= 0)
            ;
        for (int i = 0; i < n_router; ++i)
            if (i != myid && FD_ISSET(sockfd[i], &ready_set))
            {
                // if (Recv(sockfd[i], &msgr) <= 0)
                // {
                //     perror("Recv failed.");
                //     exit(EXIT_FAILURE);
                // }
                MessageR2R msgr;
                Rio_readn(sockfd[i], &msgr, msgr.len);
#ifdef _VERBOSE
                printf("[DEBUG]: read message from router %d (seq: %d, msg_type: %d, rpl_type: %d, len: %d, source: %d).\n", i, msgr.seq, msgr.message_type, msgr.reply_type, msgr.len, msgr.source);
                // printf("sockfd %d\n", i);
                assert(msgr.source == i);
                assert(msgr.reply_type == REPLY || msgr.reply_type == SEND);
#endif
                if (msgr.reply_type == SEND)
                {
#ifdef _VERBOSE
                    printf("[DEBUG]: Receive signal from router %d (seq: %d).\n", msgr.source, msgr.seq);
#endif
                    set_neighbor_dv(msgr.source, msgr.dv);
                    if (BF())
                    {
                        ++SEQ;
                        PropMsg ___(msgr.source, msgr.seq);
                        if (!___.waitlist.empty())
                            M[SEQ] = ___;
                        // wait.insert(std::make_pair(seq, new std::set<int>));

                        for (int i = 0; i < n_router; ++i)
                            if (i != myid && weight[i] < INF)
                            {
                                MessageR2R msgr2r(SEND, myid, SEQ);
                                memcpy(msgr2r.dv, dv[myid], sizeof(msgr2r.dv));
                                Rio_writen(sockfd[i], &msgr2r, msgr2r.len);
#ifdef _VERBOSE
                                printf("[DEBUG]: Send signal to router %d (seq: %d).\n", i, SEQ);
#endif
                            }
                    }
                    else
                    {
                        MessageR2R msgs(REPLY, myid, msgr.seq);
#ifdef _VERBOSE
                        printf("[DEBUG]: Send reply to router %d (src_seq: %d) // not change.\n", msgr.source, msgr.seq);
#endif
                        Rio_writen(sockfd[msgr.source], &msgs, msgs.len);
                    }
                }
                else if (msgr.reply_type == REPLY)
                {
#ifdef _VERBOSE
                    printf("[DEBUG]: Receive reply from router %d (seq: %d).\n", msgr.source, msgr.seq);
#endif
                    if (!M.count(msgr.seq) || !M[msgr.seq].del(msgr.source))
                    {
                        puts("fatal err!");
                        exit(EXIT_FAILURE);
                    }
                    if (M[msgr.seq].waitlist.empty())
                    {
                        if (M[msgr.seq].father_id != -1)
                        {
                            MessageR2R msgs(REPLY, myid, M[msgr.seq].father_seq);
#ifdef _VERBOSE
                            printf("[DEBUG]: Send reply to router %d (src_seq: %d) // message %d ok .\n", M[msgr.seq].father_id, M[msgr.seq].father_seq, msgr.seq);
#endif
                            Rio_writen(sockfd[M[msgr.seq].father_id], &msgs, msgs.len);
                        }
                        M.erase(msgr.seq);
                    }
                }
                break;
            }
    }

    // close(fd);
}

void Listen()
{
    char buf[MAXLINE];

    fd_set read_set, ready_set;
    FD_ZERO(&read_set);
    for (int i = 0; i < n_router; ++i)
        FD_SET(sockfd[i], &read_set);

    while (1)
    {

        ready_set = read_set;
        while (select(maxfd, &ready_set, NULL, NULL, NULL) <= 0)
            ;
        for (int i = 0; i < n_router; ++i)
            if (FD_ISSET(sockfd[i], &ready_set))
            {
#ifdef _VERBOSE
                printf("[DEBUG]: sockfd %d unblock.\n", i);
#endif
                Message *msg;
                Rio_readn(sockfd[i], buf, sizeof(Message));
                msg = (Message *)buf;
                Rio_readn(sockfd[i], buf + sizeof(Message), msg->len - sizeof(Message));

#ifdef _VERBOSE
                puts("");
                puts("------------------");
                if (i == myid)
                    printf("[DEBUG]: receive %d bytes from agent\n", msg->len);
                else
                    printf("[DEBUG]: receive %d bytes from %d\n", msg->len, i);
                puts("---");
#endif

                if (msg->message_type == A2R)
                {
                    MessageA2R *msga2r = (MessageA2R *)buf;
                    if (msga2r->command_type == DV)
                    {
                        puts("[INFO]: DV");

                        if (BF())
                            propogate();

                        assert(M.empty() && PropMsg::wait_cnt == 0);

                        MessageR2A msgs(0);
                        Rio_writen(sockfd[myid], &msgs, msgs.len);
#ifdef _VERBOSE
                        printf("[INFO]: send the reply of DV with payload %d bytes:\n", msgs.pl_len);
#endif
                    }
                    else if (msga2r->command_type == UPDATE)
                    {
                        MessageA2RUpdate *upd = (MessageA2RUpdate *)buf;
                        printf("[INFO]: Update: %d %d %d\n", upd->from, upd->to, upd->value);

                        weight[upd->to] = upd->value == -1 ? INF : upd->value;

                        MessageR2A msgs(0);
                        Rio_writen(sockfd[myid], &msgs, msgs.len);
#ifdef _VERBOSE
                        printf("[INFO]: send the reply of UPDATE with payload %d bytes:\n", msgs.pl_len);
#endif
                    }
                    else if (msga2r->command_type == SHOW)
                    {
                        puts("[INFO]: Show");

                        char buf[MAXLINE], str[MAXLINE];
                        get_table(str);
                        MessageR2A msgs(strlen(str) + 1);
                        // msgs.length = strlen(str) + 1;

                        memcpy(buf, &msgs, sizeof(MessageR2A));
                        memcpy(buf + sizeof(MessageR2A), str, msgs.pl_len);
                        Rio_writen(sockfd[myid], buf, msgs.len);
#ifdef _VERBOSE
                        printf("[INFO]: send the reply of SHOW with payload %d bytes:\n", msgs.pl_len);
                        puts(str);
#endif
                    }
                    else if (msga2r->command_type == RESET)
                    {
                    }
                    else
                    {
                        perror("error command type.");
                    }
                }
                else if (msg->message_type == R2R)
                {
                    MessageR2R *msgr2r = (MessageR2R *)msg;
                    if (msgr2r->reply_type == SEND)
                        set_neighbor_dv(msgr2r->source, msgr2r->dv);

                    printf("[DEBUG]: Receive signal from router %d (seq: %d).\n", msgr2r->source, msgr2r->seq);

                    if (BF())
                        propogate();

                    assert(M.empty() && PropMsg::wait_cnt == 0);

                    MessageR2R msgs(REPLY, myid, msgr2r->seq);
                    Rio_writen(sockfd[msgr2r->source], &msgs, msgs.len);
#ifdef _VERBOSE
                    printf("[DEBUG]: Send reply to router %d (src_seq: %d) // propogate finished.\n", msgr2r->source, msgr2r->seq);
#endif
                }
                break;
            }
    }
}
int main(int argc, char *argv[])
{
    setvbuf(stdout, NULL, _IONBF, 0);
    if (argc != 4)
    {
        puts("Usage: ./router <router location file> <topology conf file> <router_id>");
        return 0;
    }
    FILE *fp = fopen(argv[1], "r");
    if (fp == NULL)
    {
        puts("router location file not exist.");
        return 0;
    }

    fscanf(fp, "%d", &n_router);
    char buf[MAXLINE];

    alloc_id[-1] = -1;
    alloc_id_inv[-1] = -1;
    for (int i = 0; i < n_router; ++i)
    {
        fscanf(fp, "%s", buf);
        all_router[i].set_router(buf, i);
    }
    myid = alloc_id[atoi(argv[3])];
    this_router = &all_router[myid];
    fclose(fp);

    fp = fopen(argv[2], "r");
    if (fp == NULL)
    {
        puts("topology conf file not exist.");
        return 0;
    }

    int n_edge;
    fscanf(fp, "%d", &n_edge);
    for (int i = 0; i < n_router; ++i)
        weight[i] = INF;
    weight[myid] = 0;
    for (int i = 0; i < n_router; ++i)
    {
        for (int j = 0; j < n_router; ++j)
            dv[i][j] = INF;
        dv[i][i] = 0;
    }

    for (int i = 0; i < n_edge; ++i)
    {
        int from, to, value;
        fscanf(fp, "%d,%d,%d", &from, &to, &value);
        from = alloc_id[from];
        to = alloc_id[to];
        if (from == myid)
            weight[to] = (value == -1 ? INF : value);
    }

    for (int i = 0; i < n_router; ++i)
    {
        routing_table[i].dest = i;
        if (weight[i] < INF)
            routing_table[i].next = i,
            routing_table[i].cost = weight[i];
        else
            routing_table[i].next = -1,
            routing_table[i].cost = INF;
    }

    printf("[INFO]: router %s set up, alloc_id: %d, port: %d\n", argv[3], myid, all_router[myid].router_port);

    int listenfd = socket(AF_INET, SOCK_STREAM, 0);
    if (bind(listenfd, (struct sockaddr *)&this_router->router_addr, this_router->addr_len))
    {
        puts("bind failed");
        perror("bind failed");
        exit(EXIT_FAILURE);
    }
    printf("[INFO]: bind success\n");
    if (listen(listenfd, 1024) < 0)
    {
        puts("listen failed");
        perror("listen failed");
        exit(EXIT_FAILURE);
    }
    printf("[INFO]: listen success\n");

    while (1)
    {
        bool ok = 0;
        for (int i = 0; i <= myid && !ok; ++i)
            if (!conn[i])
                ok = 1;
        if (!ok)
            break;
        struct sockaddr_in connaddr;
        socklen_t connlen = sizeof(connaddr);
        int connfd = accept(listenfd, (struct sockaddr *)&connaddr, &connlen);
        if (connfd < 0)
        {
            perror("accept failed");
            exit(EXIT_FAILURE);
        }
        int id;
        Rio_readn(connfd, &id, sizeof(int));
        if (id == -1)
        {
            conn[myid] = 1;
            sockfd[myid] = connfd;
            printf("[INFO]: build connect to agent succesfully!\n");
        }
        else
        {
            conn[id] = 1;
            sockfd[id] = connfd;
            printf("[INFO]: build connect to router %d succesfully!\n", id);
        }
    }

    for (int i = myid + 1; i < n_router; ++i)
    {
        sockfd[i] = socket(AF_INET, SOCK_STREAM, 0);
        while (connect(sockfd[i], (struct sockaddr *)&all_router[i].router_addr, all_router[i].addr_len) < 0)
            ;
        Rio_writen(sockfd[i], &myid, sizeof(int));

        printf("[INFO]: build connect to router %d succesfully!\n", i);
        conn[i] = 1;
    }

    for (int i = 0; i < n_router; ++i)
        maxfd = std::max(maxfd, sockfd[i]);
    maxfd++;

    Listen();
    for (int i = 0; i < n_router; ++i)
        close(sockfd[i]);
    return 0;
}